#!/bin/bash

#set -x

VMIMAGE="registry.cn-hangzhou.aliyuncs.com/wxx_image/centos:dockervm"
QEMU="/opt/linux-kernel/qemu/usr/local/bin/qemu-system-x86_64"
CPUS=2
MEMORY=4096M
QEMU="/usr/libexec/qemu-kvm"
QEMU_NET_SCRIPT="/var/dockervm/qemu-ifup"
QEMU_NET_DOWNSCRIPT="/var/dockervm/qemu-ifdown"
QEMU_NET_MAC=$(echo 52:`od /dev/urandom -w5 -tx1 -An|head -n 1|sed -e 's/ //' -e 's/ /:/g'`)
QEMU_SOCK="/run/dockervm.sock"
GUEST_OS=""

QEMU_OPT="-smp $CPUS -m $MEMORY -nographic -enable-kvm -drive file=/var/dockervm/linux.img,if=none,id=drive-virtio-disk0 -device virtio-blk-pci,drive=drive-virtio-disk0,id=virtio-disk0 -serial unix:$QEMU_SOCK,server,nowait"

CONTAINERID=""
CONTAINER_NAME=""
NET_TYPE="user"

IFNAME=br0
CONTAINER_IFNAME=eth0
FAMILY_FLAG="-4"

function try_exec() {
        $@
        if [ $? -ne 0 ]
        then
                exit 1
        fi
}

function create_container_network() {
	if [ -z "$1" ]; then
    	return 1
  	fi
  	local cid=$1
  	GATEWAY=`route -n | grep $IFNAME | grep UG | awk '{print $2}'`
        if [ $? -ne 0 ]
        then
                return 1
        fi

        DOCKERPID=`docker inspect --format='{{ .State.Pid }}' "$cid"`
        if [ $? -ne 0 ]
        then
                return 1
        fi
        DOCKERCID=`docker inspect --format='{{ .ID }}' "$cid"`
        if [ $? -ne 0 ]
        then
          return 1
        fi

        DOCKERCNAME=`docker inspect --format='{{ .Name }}' "$cid"`
        if [ $? -ne 0 ]
    then
                return 1
        fi


        NSPID=$DOCKERPID

        LOCAL_IFNAME="v${CONTAINER_IFNAME}pl${NSPID}"
        GUEST_IFNAME="v${CONTAINER_IFNAME}pg${NSPID}"
        MTU=`ip link show "$IFNAME" | awk '{print $5}'`
        if [ $? -ne 0 ]
    then
        return 1
    fi

        #
        try_exec mkdir -p /var/run/netns
        try_exec rm -f "/var/run/netns/$NSPID"
        try_exec ln -s "/proc/$NSPID/ns/net" "/var/run/netns/$NSPID"

        # create veth-pair
        try_exec ip link add name "$LOCAL_IFNAME" mtu "$MTU" type veth peer name "$GUEST_IFNAME" mtu "$MTU"
        try_exec ip link set "$LOCAL_IFNAME" master "$IFNAME" > /dev/null 2>&1
        try_exec ip link set "$LOCAL_IFNAME" up

        # set veth-pair with container namespace
        try_exec ip link set "$GUEST_IFNAME" netns "$NSPID"
        try_exec ip netns exec "$NSPID" ip link set "$GUEST_IFNAME" name "$CONTAINER_IFNAME"
        # Remove NSPID to avoid `ip netns` catch it.
        try_exec rm -f "/var/run/netns/$NSPID"
}

function parse_create_arg() {
	while getopts ":N:n:I:" OPTION
	do
		case $OPTION in
			N ) NET_TYPE="$OPTARG";;
			n ) CONTAINER_NAME="$OPTARG";;
			I ) GUEST_OS="$OPTARG";;
		esac
	done
	if [ "$NET_TYPE" = "bridge" ]; then
		QEMU_OPT="$QEMU_OPT -netdev tap,id=hostnet0,script=$QEMU_NET_SCRIPT,downscript=$QEMU_NET_DOWNSCRIPT -device virtio-net-pci,netdev=hostnet0,id=net0,mac=$QEMU_NET_MAC"
	else
		QEMU_OPT="$QEMU_OPT -net nic -net user,hostfwd=tcp::10022-:22"
	fi

	if [ -n "$GUEST_OS" ]; then
		VMIMAGE="$VMIMAGE"_"$GUEST_OS"
	fi
}


function create() {
	local name_opt=""
	local net_opt=""

	if [ -n "$CONTAINER_NAME" ]; then
		local name_opt="--name $CONTAINER_NAME"
	fi
	if [ "$NET_TYPE" = "bridge" ]; then
		local net_opt="--network none"
	fi
	local cid=$(docker run -td $name_opt --cap-add NET_ADMIN  --device /dev/net/tun:/dev/net/tun:rw \
		--device /dev/kvm:/dev/kvm:rw --env QEMU=$QEMU --env QEMU_OPT="$QEMU_OPT" $net_opt \
		--env NET_TYPE=$NET_TYPE --tmpfs /run $VMIMAGE /var/dockervm/start_qemu.sh)
	if [ "$NET_TYPE" = "bridge" ]; then
		try_exec create_container_network $cid
	fi
	try_exec docker exec $cid touch /run/dockervm_net.done
}

function parse_start_arg() {
        CONTAINERID=$1
        if [ -z "$CONTAINERID" ]; then
                echo "Invalid container id"
                exit 1
        fi
}

function start() {
	local net_type_env=$(docker inspect -f '{{range $index, $value := .Config.Env}}{{println $value}}{{end}}' $CONTAINERID | grep NET_TYPE)
	local net_type=$(echo $net_type_env | cut -d = -f 2)
	stat=`docker inspect --format {{.State.Status}} $CONTAINERID`
	if [ "$stat" != "running" ]; then
		try_exec docker start $CONTAINERID
		if [ "$net_type" = "bridge" ]; then
			try_exec create_container_network $CONTAINERID
		fi
		try_exec docker exec $CONTAINERID touch /run/dockervm_net.done
	fi
}

function parse_attach_arg() {
	CONTAINERID=$1
	if [ -z "$CONTAINERID" ]; then
		echo "Invalid container id"
		exit 1
	fi
}

# exit with "control + q"
function attach() {
	docker exec -it $CONTAINERID socat "stdin,raw,echo=0,escape=0x11" "unix-connect:${QEMU_SOCK}"
}


function parse_delete_arg() {
    CONTAINERID=$1
    if [ -z "$CONTAINERID" ]; then
        echo "Invalid container id"
        exit 1
    fi
}

function delete() {
	docker rm -f $CONTAINERID
}

function usage() {
	echo ""
	echo "Usage:  dockervm COMMAND"
	echo ""
	echo "A tool to run vm in docker container"
	echo ""
	echo "  dockervm create [-N user(default)|bridge -I OS -n NAME]"
	echo "  dockervm start  CONTAINER"
	echo "  dockervm attach CONTAINER"
	echo "  dockervm rm     CONTAINER"
}

function main() {
	if [ $# -lt 1 ]; then
		usage
		exit 0
	fi
	OPTYPE=$1

	if [ $OPTYPE = "create" ]; then
		shift 1
		parse_create_arg $*
		create
	elif [ $OPTYPE = "start" ]; then
		shift 1
		parse_start_arg $*
		start
	elif [ $OPTYPE = "attach" ]; then
		shift 1
		parse_attach_arg $*
		attach
	elif [ $OPTYPE = "rm" ]; then
		shift 1
		parse_delete_arg $*
		delete
	else
		echo "Invalid command"
		usage
		exit 1
	fi
}

main $@
